"""
ref: 
    https://github.com/pnnx/pnnx
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple pnnx
"""
import os
import numpy as np
import torch
import torch.nn as nn 
import torch.nn.functional as F
import struct
import onnxruntime as rt


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = nn.Linear(512, 1000)

    def forward(self, x):
        x = self.linear(x)
        return x


def export_onnx(model, onnx_path):
    inputs = torch.randn(1, 512)
    torch.onnx.export(model, inputs, onnx_path, 
        verbose=True,
        opset_version=11, 
        input_names=['x'], 
        output_names=['y'], 
    )


def infer(model):
    x = np.array([0.0, 0.1, 0.2], dtype=np.float32)
    x = torch.from_numpy(x)
    y = model(x)
    print(y)


def infer_onnx(onnx_path):
    x = np.array([0.0, 0.1, 0.2], dtype=np.float32)
    x = np.expand_dims(x, axis=0)
    session = rt.InferenceSession(onnx_path)
    input_map = {'x': x}
    output_names = ['y']
    output = session.run(output_names, input_map)
    print(output[0])


if __name__ == '__main__':
    model = Model()
    # infer(model)
    export_onnx(model, "./linear_large.onnx")
    # infer_onnx('../data/linear.onnx')
